# 行分割保存到excel里且支持追加保存
# 支持保存为txt格式
# 增加“帮助”按钮
# 优化界面
# 改变界面
# 增加批量识别及保存
# 支持批量识别并保存为excel
# 大字段名动态生成，而非自己预先定义。
# 保存为txt文件时候顺便保存时间
# 修改保存路径为当前文件夹
from PyQt6.QtWidgets import (QApplication, QMessageBox, QInputDialog, QFileDialog,
                             QProgressDialog)
from PyQt6.QtGui import QPixmap, QDragEnterEvent, QDropEvent
from PyQt6.QtCore import Qt
import cv2
from paddleocr import PaddleOCR
import os
from ui import BaseApp
from io import StringIO
import sys
from save_results import SaveResults

try:
    import pyi_splash
    pyi_splash.close()
except ImportError:
    pass

class MyApp(BaseApp):
    def __init__(self):
        super().__init__()
        self.ocr_engine = None
        self.current_image_path = None
        self.output_folder = os.getcwd()  # 获取当前工作目录
        self.excel_file_name = "OCR_Results.xlsx"
        self.processing = False  # 防止重复点击
        # 新增属性
        self.image_paths = []  # 存储所有图片路径
        self.current_results = {}  # {图片路径: [原始文本行1, 原始文本行2, ...]}

        # 初始化OCR引擎
        self.init_ocr_engine()

        # 连接信号
        self.upload_btn.clicked.connect(self.open_image_dialog)
        self.ocr_btn.clicked.connect(self.run_ocr)
        self.select_folder_btn.clicked.connect(self.select_output_folder)
        self.input_filename_btn.clicked.connect(self.input_file_name)
        self.save_excel_btn.clicked.connect(self.save_results)
        self.lang_combo.currentIndexChanged.connect(self.init_ocr_engine)
        # self.help_btn.clicked.connect(self.show_help)
        self.setAcceptDrops(True)
        self.setWindowTitle("OCR文字识别工具")
        # 连接新按钮信号
        self.add_folder_btn.clicked.connect(self.add_image_folder)
        self.clear_btn.clicked.connect(self.clear_image_list)
        self.image_list.itemClicked.connect(self.on_image_selected)
        # 初始化保存结果的类
        self.save_results_handler = SaveResults(self)

        # 替换 save_results 方法

    def save_results(self):
        self.save_results_handler.save_results()

    def init_ocr_engine(self):
        lang = 'ch' if self.lang_combo.currentText() == "中文" else 'en'
        try:
            self.ocr_engine = PaddleOCR(use_angle_cls=True, lang=lang, use_gpu=False)
        except Exception as e:
            QMessageBox.critical(self, "引擎初始化失败", f"无法初始化OCR引擎：{str(e)}")

    def open_image_dialog(self):
        files, _ = QFileDialog.getOpenFileNames(
            self, "选择图片", "",
            "图片文件 (*.png *.jpg *.jpeg *.bmp)"
        )
        if files:
            for path in files:
                if path not in self.image_paths:
                    self.image_paths.append(path)
                    self.image_list.addItem(path)

    # 新增方法：清空列表
    def clear_image_list(self):
        self.image_paths.clear()
        self.image_list.clear()
        self.image_label.clear()
        self.current_image_path = None

    # 实现图片预览
    def on_image_selected(self, item):
        path = item.text()
        self.current_image_path = path
        self.display_image(path)

    def display_image(self, path):
        self.current_image_path = path
        pixmap = QPixmap(path)
        if not pixmap.isNull():
            scaled_pixmap = pixmap.scaled(
                self.image_label.width(),
                self.image_label.height(),
                Qt.AspectRatioMode.KeepAspectRatio,
                Qt.TransformationMode.SmoothTransformation
            )
            self.image_label.setPixmap(scaled_pixmap)
        else:
            self.show_error("图片加载失败", "不支持的文件格式或损坏的图片文件")

    def preprocess_image(self, image_path):
        try:
            img = cv2.imread(image_path)
            if img is None:
                raise ValueError("无法读取图片文件")

            # 显示预处理进度
            progress = QProgressDialog("正在预处理图片...", None, 0, 0, self)
            progress.setWindowModality(Qt.WindowModality.WindowModal)
            progress.show()
            QApplication.processEvents()  # 更新UI

            # ========== 改进的预处理流程 ========== #
            # 1. 转换为灰度图
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

            # 2. 使用CLAHE增强对比度（来自test8.py）
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            enhanced = clahe.apply(gray)

            # 3. 自适应阈值二值化（替换原有方法）
            _, thresh = cv2.threshold(enhanced, 150, 255, cv2.THRESH_BINARY)

            # 生成临时文件路径
            temp_path = os.path.join(os.path.dirname(__file__), "temp_preprocessed.png")
            cv2.imwrite(temp_path, thresh)
            # ==================================== #

            progress.close()
            return temp_path

        except Exception as e:
            self.show_error("预处理失败", str(e))
            return None

    # 新增方法：添加文件夹
    def add_image_folder(self):
        folder = QFileDialog.getExistingDirectory(self, "选择图片文件夹")
        if folder:
            self.find_images_in_folder(folder)

    # 新增方法：递归查找图片
    def find_images_in_folder(self, folder):
        extensions = ('.png', '.jpg', '.jpeg', '.bmp')
        for root, _, files in os.walk(folder):
            for file in files:
                if file.lower().endswith(extensions):
                    path = os.path.join(root, file)
                    if path not in self.image_paths:
                        self.image_paths.append(path)
                        self.image_list.addItem(path)

    # 修改OCR执行方法
    def run_ocr(self):
        if not self.image_paths:
            self.show_warning("提示", "请先添加要识别的图片")
            return

        self.is_cancelled = False  # 取消标志位

        self.result_area.clear()
        self.current_results.clear()

        # 创建进度对话框
        self.progress = QProgressDialog("正在识别图片...", "取消", 0, len(self.image_paths), self)
        self.progress.setWindowModality(Qt.WindowModality.WindowModal)
        self.progress.canceled.connect(self.on_cancel_ocr)  # 连接取消信号

        try:
            for i, path in enumerate(self.image_paths):
                self.progress.setValue(i)
                QApplication.processEvents()  # 处理事件循环，确保UI响应

                if self.is_cancelled:  # 如果用户点击了取消按钮，退出循环
                    break

                try:
                    processed_path = self.preprocess_image(path)
                    result = self.ocr_engine.ocr(processed_path, cls=True)

                    # 处理空结果
                    if not result or not result[0]:
                        error_msg = f"⚠️ 未识别到文字：{os.path.basename(path)}"
                        self.result_area.append(error_msg)
                        self.result_area.append("-" * 50 + "\n")
                        continue

                    raw_lines = [line[1][0].strip() for line in result[0]]  # 提取原始文本行
                    self.current_results[path] = raw_lines  # 存储为列表而非拼接字符串

                    # 展示结果时仍可拼接显示（不影响保存逻辑）
                    display_text = "\n".join(raw_lines)
                    self.result_area.append(f"✅ 识别成功：{os.path.basename(path)}")
                    self.result_area.append(display_text)
                    self.result_area.append("=" * 50 + "\n")

                except Exception as e:
                    # 修改此处：将错误显示在结果区域而不是弹窗
                    error_msg = f"""
        ❌ 识别失败：{os.path.basename(path)}
        错误类型：{type(e).__name__}
        错误详情：{str(e)}
                    """
                    self.result_area.append(error_msg)
                    self.result_area.append("-" * 50 + "\n")

                finally:
                    if 'processed_path' in locals() and os.path.exists(processed_path):
                        os.remove(processed_path)

        finally:
            # 显式断开信号并关闭对话框
            if hasattr(self, 'progress'):
                self.progress.canceled.disconnect()  # 断开所有信号
                self.progress.close()

        # 根据是否取消显示不同的提示
        if self.is_cancelled:
            self.result_area.append("\n▶ 识别任务已取消")
            self.show_info("取消", "识别任务已取消")
        else:
            self.show_info("完成", "所有图片识别完成！")
            self.result_area.append("\n▶ 识别任务已完成")

    def on_cancel_ocr(self):
        """当用户点击取消按钮时执行的操作"""
        if not self.is_cancelled:
            self.is_cancelled = True
            if hasattr(self, 'progress'):
                # 更新UI但不关闭对话框（由run_ocr统一关闭）
                self.progress.setLabelText("正在取消...")
                self.progress.setCancelButton(None)  # 移除取消按钮

    # 修改后的拖放事件处理
    def dragEnterEvent(self, event: QDragEnterEvent):
        if event.mimeData().hasUrls():
            event.acceptProposedAction()

    def dropEvent(self, event: QDropEvent):
        if self.processing:
            return

        # 获取所有拖放路径并过滤图片文件
        for url in event.mimeData().urls():
            path = url.toLocalFile()
            if os.path.isfile(path) and path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
                if path not in self.image_paths:
                    self.image_paths.append(path)
                    self.image_list.addItem(path)
            elif os.path.isdir(path):  # 处理拖放文件夹
                self.find_images_in_folder(path)

        # 显示第一个图片预览
        if self.image_paths:
            self.display_image(self.image_paths[0])

    def select_output_folder(self):
        folder_path = QFileDialog.getExistingDirectory(self, "选择保存位置", self.output_folder)
        if folder_path:
            self.output_folder = folder_path
            self.save_results_handler.update_output_folder(folder_path)  # 更新 SaveResults 中的保存路径
            QMessageBox.information(self, "路径已更新", f"文件将保存到：\n{folder_path}")

    def input_file_name(self):
        # 去掉时间戳逻辑，使用固定基础名称
        default_name = "OCR_Results"  # 固定基础名称
        text, ok = QInputDialog.getText(
            self, "输入文件名",
            "请输入文件名（不包含扩展名）：",
            text=default_name  # 默认显示固定名称
        )
        # 如果用户输入为空则使用默认名称
        self.excel_file_name = f"{text.strip() or default_name}.xlsx"

    # 辅助提示方法
    def show_error(self, title, message):
        QMessageBox.critical(self, title, message)

    def show_warning(self, title, message):
        QMessageBox.warning(self, title, message)

    def show_info(self, title, message):
        QMessageBox.information(self, title, message)

    if sys.stderr is None:
        sys.stderr = StringIO()

if __name__ == "__main__":
    import sys
    app = QApplication(sys.argv)
    app.setStyle("Fusion")  # 使用更现代的主题
    window = MyApp()
    window.show()
    sys.exit(app.exec())
